{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Open-Loop Evaluation\n",
    "\n",
    "In this notebook you are going to evaluate a CNN-based policy to control the SDV with a protocol named *open-loop* evaluation.\n",
    "\n",
    "**Note: this notebook assumes you've already run the [training notebook](./train.ipynb) and stored your model successfully.**\n",
    "\n",
    "## What is open-loop evaluation?\n",
    "In open-loop evaluation we evaluate our model prediction as we follow the annotated ground truth.\n",
    "\n",
    "In each frame, we compare the predictions of our model against the annotated ground truth. This can be done with different metrics, and we will see a few of them in the following.\n",
    "\n",
    "**Regardless of the metric used, this evaluation protocol doesn't modify the future locations according to our predictions.**\n",
    "\n",
    "![open-loop](../../docs/images/planning/open-loop.svg)\n",
    "\n",
    "\n",
    "## What can we use open-loop evaluation for?\n",
    "Open-loop evaluation can be used for a frame by frame comparison between the expert and the policy. This is extremely useful for debugging the model behaviours and investigate outlier predictions in specific situations (e.g. at crossings or unprotected turn).\n",
    "\n",
    "## Is open-loop evaluation enough?\n",
    "Regardless of the quality of the open-loop results, **this evaluation is not enough** to ensure your model will be able to actually drive on the road (that's where we all want to go in the end). If your model is not in full control of the SDV, you can't really say it will work once the annotated trajectory won't be available anymore.\n",
    "\n",
    "Before drawing conclusions on our model we must test it when it is in full control of the SDV, in a setting called **closed-loop**. You can try just that in our [dedicated closed-loop evaluation notebook](./closed_loop_test.ipynb)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from tempfile import gettempdir\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.dataloader import default_collate\n",
    "from tqdm import tqdm\n",
    "\n",
    "from l5kit.configs import load_config_data\n",
    "from l5kit.data import LocalDataManager, ChunkedDataset\n",
    "from l5kit.dataset import EgoDataset\n",
    "from l5kit.rasterization import build_rasterizer\n",
    "from l5kit.geometry import transform_points, angular_distance\n",
    "from l5kit.visualization import TARGET_POINTS_COLOR, PREDICTED_POINTS_COLOR, draw_trajectory\n",
    "from l5kit.kinematic import AckermanPerturbation\n",
    "from l5kit.random import GaussianRandomGenerator\n",
    "\n",
    "import os"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Prepare data path and load cfg\n",
    "\n",
    "By setting the `L5KIT_DATA_FOLDER` variable, we can point the script to the folder where the data lies.\n",
    "\n",
    "Then, we load our config file with relative paths and other configurations (rasteriser, training params...)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# set env variable for data\n",
    "os.environ[\"L5KIT_DATA_FOLDER\"] = \"/tmp/l5kit_data\"\n",
    "dm = LocalDataManager(None)\n",
    "# get config\n",
    "cfg = load_config_data(\"./config.yaml\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load the model"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "model_path = \"/tmp/planning_model.pt\"\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = torch.load(model_path).to(device)\n",
    "model = model.eval()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load the evaluation data\n",
    "This is almost the exact same code you've already seen in the [training notebook](./train.ipynb). Apart from the different dataset we load, the biggest difference is that **we don't perturb our data here**.\n",
    "\n",
    "When performing evaluation we're interested in knowing the performance on the annotated data, not on perturbed one."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# ===== INIT DATASET\n",
    "eval_cfg = cfg[\"val_data_loader\"]\n",
    "rasterizer = build_rasterizer(cfg, dm)\n",
    "eval_zarr = ChunkedDataset(dm.require(eval_cfg[\"key\"])).open()\n",
    "eval_dataset = EgoDataset(cfg, eval_zarr, rasterizer)\n",
    "eval_dataloader = DataLoader(eval_dataset, shuffle=eval_cfg[\"shuffle\"], batch_size=eval_cfg[\"batch_size\"], \n",
    "                             num_workers=eval_cfg[\"num_workers\"])\n",
    "print(eval_dataset)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Evaluation loop\n",
    "\n",
    "Here, we loop through the data and store predicted and annotated trajectories (positions + yaws).\n",
    "\n",
    "Note: we're not taking into account availability here. We acknowledge this can reflect in a lower score."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# ==== EVAL LOOP\n",
    "position_preds = []\n",
    "yaw_preds = []\n",
    "\n",
    "position_gts = []\n",
    "yaw_gts = []\n",
    "\n",
    "torch.set_grad_enabled(False)\n",
    "\n",
    "for idx_data, data in enumerate(tqdm(eval_dataloader)):\n",
    "    data = {k: v.to(device) for k, v in data.items()}\n",
    "    result = model(data)\n",
    "    position_preds.append(result[\"positions\"].detach().cpu().numpy())\n",
    "    yaw_preds.append(result[\"yaws\"].detach().cpu().numpy())\n",
    "\n",
    "    position_gts.append(data[\"target_positions\"].detach().cpu().numpy())\n",
    "    yaw_gts.append(data[\"target_yaws\"].detach().cpu().numpy())\n",
    "    if idx_data == 10:\n",
    "        break\n",
    "    \n",
    "position_preds = np.concatenate(position_preds)\n",
    "yaw_preds = np.concatenate(yaw_preds)\n",
    "\n",
    "position_gts = np.concatenate(position_gts)\n",
    "yaw_gts = np.concatenate(yaw_gts)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Quantitative evaluation"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## ADE, FDE and angle distance\n",
    "\n",
    "### Positional displacement\n",
    "[Average Displacement Error (ADE) and Final Displacement Error (FDE)](https://en.wikipedia.org/wiki/Mean_squared_displacement) are standard metrics used to evaluate future predictions for AVs.\n",
    "\n",
    "We can compute them by comparing predicted and annotated positions, which we have stored in the previous cell.\n",
    "Additionally, we can plot histograms of their distributions across samples to better capture the variance of our error.\n",
    "\n",
    "### Angle displacement\n",
    "\n",
    "For the yaw, we can use the Minimum Angle Distance to check the error. Again, we can plot a histogram to inspect the error distribution. \n",
    "\n",
    "Although yaw may seem redundant here, it's actually crucial to fully control the SDV. We'll use it extensively in the closed-loop evaluation notebook."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "pos_errors = np.linalg.norm(position_preds - position_gts, axis=-1)\n",
    "\n",
    "# DISPLACEMENT AT T\n",
    "plt.plot(np.arange(pos_errors.shape[1]), pos_errors.mean(0), label=\"Displacement error at T\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# ADE HIST\n",
    "plt.hist(pos_errors.mean(-1), bins=100, label=\"ADE Histogram\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# FDE HIST\n",
    "plt.hist(pos_errors[:,-1], bins=100, label=\"FDE Histogram\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "angle_errors = angular_distance(yaw_preds, yaw_gts).squeeze()\n",
    "\n",
    "# ANGLE ERROR AT T\n",
    "plt.plot(np.arange(angle_errors.shape[1]), angle_errors.mean(0), label=\"Angle error at T\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# ANGLE ERROR HIST\n",
    "plt.hist(angle_errors.mean(-1), bins=100, label=\"Angle Error Histogram\")\n",
    "plt.legend()\n",
    "plt.show()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Qualitative evaluation"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Visualise results\n",
    "We can also visualise some images with predicted and annotated trajectories using L5Kit visualisation features.\n",
    "\n",
    "In this example, we draw 20 images from our dataset and we visualise predicted and annotated trajectories on top of them."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "for frame_number in range(0, len(eval_dataset), len(eval_dataset) // 20):\n",
    "    \n",
    "    data = eval_dataloader.dataset[frame_number]\n",
    "    data_batch = default_collate([data])\n",
    "    data_batch = {k: v.to(device) for k, v in data_batch.items()}\n",
    "    result = model(data_batch)\n",
    "    predicted_positions = result[\"positions\"].detach().cpu().numpy().squeeze()\n",
    "\n",
    "    im_ego = rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0))\n",
    "    target_positions = data[\"target_positions\"]\n",
    "    \n",
    "    predicted_positions = transform_points(predicted_positions, data[\"raster_from_agent\"])\n",
    "    target_positions = transform_points(target_positions, data[\"raster_from_agent\"])\n",
    "    \n",
    "    draw_trajectory(im_ego, predicted_positions, PREDICTED_POINTS_COLOR)\n",
    "    \n",
    "    draw_trajectory(im_ego, target_positions, TARGET_POINTS_COLOR)\n",
    "\n",
    "    plt.imshow(im_ego)\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Visualise the open-loop\n",
    "\n",
    "To visualise the open loop we can just repeat the same operations for the consecutive frames.\n",
    "\n",
    "In this example, we show the first 200 frames for our dataset, plotting predicted and annotated trajectories.\n",
    "\n",
    "**We want to stress this out again: this is an open loop evaluation, we are NOT controlling the AV with our predictions**"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from IPython.display import display, clear_output\n",
    "import PIL\n",
    " \n",
    "for frame_number in range(200):\n",
    "    \n",
    "    data = eval_dataloader.dataset[frame_number]\n",
    "\n",
    "    data_batch = default_collate([data])\n",
    "    data_batch = {k: v.to(device) for k, v in data_batch.items()}\n",
    "    \n",
    "    result = model(data_batch)\n",
    "    predicted_positions = result[\"positions\"].detach().cpu().numpy().squeeze()\n",
    "\n",
    "    \n",
    "    predicted_positions = transform_points(predicted_positions, data[\"raster_from_agent\"])\n",
    "    target_positions = transform_points(data[\"target_positions\"], data[\"raster_from_agent\"])\n",
    "    \n",
    "    im_ego = rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0))\n",
    "    draw_trajectory(im_ego, target_positions, TARGET_POINTS_COLOR)\n",
    "    draw_trajectory(im_ego, predicted_positions, PREDICTED_POINTS_COLOR)\n",
    "    \n",
    "    clear_output(wait=True)\n",
    "    display(PIL.Image.fromarray(im_ego))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Pre-trained model results\n",
    "\n",
    "We include here the open-loop results of one scene using one of our pre-trained model. The predicted trajectory is well overlapped with the annotated one.\n",
    "\n",
    "![SegmentLocal](../../docs/images/planning/out_9_open.gif \"segment\")"
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}